package com.jack.leetcode.eachproblem_2020;

import java.util.*;

public class 树中距离之和 {

    public static void main(String[] args) {
        int[][] edges = {{0,1},{0,2},{2,3},{2,4},{2,5}};
        sumOfDistancesInTree(6, edges);
        int[][] a = {{2,0},{1,0}};
        sumOfDistancesInTree(3, a);
        int[][] b = {{1,2},{3,2},{3,0}};
        sumOfDistancesInTree(4, b);
    }

    private static int[] dp ;
    private static int[] ans ;
    private static int[] sz ;
    private static List<List<Integer>> graph;
    public static int[] sumOfDistancesInTree(int N, int[][] edges) {
        dp = new int[N];
        ans = new int[N];
        sz = new int[N];
        graph = new ArrayList<>(2*N);
        for(int i = 0;i < N;i++){
            graph.add(new ArrayList<>());
        }
        for (int[] edge : edges) {
            graph.get(edge[0]).add(edge[1]);
            graph.get(edge[1]).add(edge[0]);
        }
        dfs1(0, -1);
        dfs2(0, -1);
        return ans;
    }

    private static void dfs1(int u, int f){
        sz[u] = 1;
        dp[u] = 0;
        for(int v : graph.get(u)){
            if(v == f){
                continue;
            }
            dfs1(v, u);
            dp[u] += dp[v] + sz[v];
            sz[u] += sz[v];
        }
    }

    private static void dfs2(int u, int f){
        ans[u] = dp[u];
        for(int v : graph.get(u)){
            if(v == f){
                continue;
            }
            int du = dp[u], dv = dp[v], su = sz[u], sv = sz[v];

            dp[u] -= dp[v] + sz[v];
            sz[u] -= sz[v];
            dp[v] += dp[u] + sz[u];
            sz[v] += sz[u];

            dfs2(v, u);

            dp[u] = du;
            dp[v] = dv;
            sz[u] = su;
            sz[v] = sv;
        }
    }
}
